package com.study.hive.udaf;

import org.apache.hadoop.hive.ql.metadata.HiveException;
import org.apache.hadoop.hive.ql.udf.generic.GenericUDAFEvaluator;
import org.apache.hadoop.hive.ql.util.JavaDataModel;
import org.apache.hadoop.hive.serde2.objectinspector.ObjectInspector;
import org.apache.hadoop.hive.serde2.objectinspector.primitive.LongObjectInspector;
import org.apache.hadoop.hive.serde2.objectinspector.primitive.PrimitiveObjectInspectorFactory;
import org.apache.hadoop.io.LongWritable;

/**
 * GenericUDAFEvaluator类实现UDAF的逻辑
 * <p>
 * enum Mode 运行阶段枚举类
 * PARTIAL1:
 * 这个是mapreduce的map阶段:从原始数据到部分数据聚合
 * 将会调用iterate()和terminatePartial()
 * PARTIAL2:
 * 这个是mapreduce的map端的Combiner阶段，负责在map端合并map的数据:部分数据聚合
 * 将会调用merge() 和 terminatePartial()
 * FINAL:
 * mapreduce的reduce阶段:从部分数据的聚合到完全聚合
 * 将会调用merge()和terminate()
 * COMPLETE:
 * 如果出现了这个阶段，表示mapreduce只有map，没有reduce，所以map端就直接出结果了:从原始数据直接到完全聚合
 * 将会调用 iterate()和terminate()
 */
public class CountUDAFEvaluator extends GenericUDAFEvaluator {
    private boolean isAllColumns = false;

    private LongObjectInspector aggOI; // 合并结果的类型

    private LongWritable result;

    public CountUDAFEvaluator() {
    }

    public CountUDAFEvaluator(boolean isAllCols) {
        isAllColumns = isAllCols;
    }

    /**
     * 初始化
     *
     * @param m          代表此时在map-reduce哪个阶段，因为不同的阶段可能在不同的机器上执行，需要重新创建对象
     *                   partial1、partial2、final、complete
     * @param parameters partial1或complete阶段传入的parameters类型是原始输入数据的类型
     *                   partial2和final阶段（执行合并）的parameters类型是partial-aggregations（既合并返回结果的类型），此时parameters长度肯定只有1了
     * @return ObjectInspector
     * 在partial1和partial2阶段返回局部合并结果的类型，既terminatePartial的类型
     * 在complete或final阶段返回总结果的类型，既terminate的类型
     */
    @Override
    public ObjectInspector init(Mode m, ObjectInspector[] parameters) throws HiveException {
        super.init(m, parameters);
        // 当是combiner和reduce阶段时，获取合并结果的类型，因为需要执行merge方法
        // merge方法需要部分合并的结果类型来取得值
        if (m == Mode.PARTIAL2 || m == Mode.FINAL) {
            aggOI = (LongObjectInspector) parameters[0];
        }
        result = new LongWritable(0); // 保存总结果
        return PrimitiveObjectInspectorFactory.writableLongObjectInspector; // 局部合并结果的类型和总合并结果的类型都是long
    }

    // 定义一个AbstractAggregationBuffer类来缓存合并值
    static class CountAgg extends AbstractAggregationBuffer {
        long value;

        // 返回类型占的字节数，long为8
        @Override
        public int estimate() {
            return JavaDataModel.PRIMITIVES2;
        }
    }

    // 创建缓存合并值的buffer
    @Override
    public AggregationBuffer getNewAggregationBuffer() throws HiveException {
        CountAgg countAgg = new CountAgg();
        reset(countAgg);
        return countAgg;
    }

    // 重置合并值
    @Override
    public void reset(AggregationBuffer agg) throws HiveException {
        ((CountAgg) agg).value = 0;
    }

    // map时执行，迭代数据
    @Override
    public void iterate(AggregationBuffer agg, Object[] parameters) throws HiveException {
        // parameters为输入数据
        // parameters == null means the input table/split is empty
        if (parameters == null) {
            return;
        }
        if (isAllColumns) {
            ((CountAgg) agg).value++;
        } else {
            boolean countThisRow = true;
            for (Object nextParam : parameters) {
                if (nextParam == null) {
                    countThisRow = false;
                    break;
                }
            }
            if (countThisRow) {
                ((CountAgg) agg).value++;
            }
        }
    }

    // 返回buffer中部分聚合结果，map结束和combiner结束执行
    @Override
    public Object terminatePartial(AggregationBuffer agg) throws HiveException {
        return terminate(agg);
    }

    // 合并结果，combiner或reduce时执行
    @Override
    public void merge(AggregationBuffer agg, Object partial) throws HiveException {
        if (partial != null) {
            ((CountAgg) agg).value += aggOI.get(partial); // 累加部分聚合的结果
        }
    }

    // 返回buffer中总结果，reduce结束执行或者没有reduce时map结束执行
    @Override
    public Object terminate(AggregationBuffer agg) throws HiveException {
        result.set(((CountAgg) agg).value); // 每一组执行一次（group by）
        return result; // 返回writable类型
    }
}
